package cc.jt.tl;

import java.time.Duration;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Supplier;

public class MyTimeLimiter {

  private final Duration timeoutDuration;
  private final boolean cancelRunningFuture;

  public MyTimeLimiter(Duration timeoutDuration) {
    this.timeoutDuration = timeoutDuration;
    this.cancelRunningFuture = false;
  }

  public MyTimeLimiter(Duration timeoutDuration, boolean cancelRunningFuture) {
    this.timeoutDuration = timeoutDuration;
    this.cancelRunningFuture = cancelRunningFuture;
  }

  public <T, F extends Future<T>> T executeFutureSupplier(Supplier<F> futureSupplier) throws Exception {
    F f = futureSupplier.get();
    try {
      return f.get(timeoutDuration.toMillis(), TimeUnit.MILLISECONDS);
    }
    catch (TimeoutException timeoutException) {
      if (cancelRunningFuture) {
        f.cancel(true);
      }
      throw timeoutException;
    }
  }
}
